library(mlogit)
library(plyr)
library(reshape2)
#library(lme4)
library(rbenchmark)
library(ggplot2)
library(scales)
library(grid)
library(Matrix)
library(grplasso)

setwd("~/Dropbox/MNL FML")
#source("R/hlogit.R")
#load("R/da.Rdata")
dyn.load("src/glasso.so")

norm <- function(x) sqrt(sum(x^2))
th_glasso <- function(x, lam)
  max(1 - lam/norm(x), 0) * x



#################################
# Step 1: Simulate data
#################################
set.seed(10)
data(splice)
da <- splice[, -(1:4)]
n <- nrow(da) 
x1 <- rnorm(n, 2)
x2 <- rnorm(n)
X <- model.matrix(~., da)[, -1]
X <- cbind(X, x1, x2)
grp <- c(0, 3, 6, 9, 12, 13, 14)
ncols <- diff(grp)

XX <- as(X, "sparseMatrix")
scales <- scaleX(XX)
X <- scales$X
alpha <- 1
beta <- c(rep(0, 6), 
          0.5, 1, -0.5, 0.5, 1, -0.5,
          0, 1)
y <- rnorm(n, as.numeric(X %*% beta + alpha))

# get lambda
maxLam <- .Call("R_mnl_get_maxLam", X, y, as.integer(grp))
lambda <- seq(maxLam, 0.001, length = 20)


#==================================#
# Step 2: Run group lasso
#==================================#

# a) C implementation

niter = 200
intercept = TRUE
trace = TRUE
#lambda = 1.1
dims <- as.integer(c(nrow(X), ncol(X), 0, 0,  
                     length(grp) - 1, length(lambda)))
names(dims) <- c("nO", "nB", "nP", "nC", "nG", "nL")

ans <- new("glasso", X = X, y = y,
           dims = dims,
           grp = as.integer(grp),
           ncols = as.integer(ncols),
           lambda = lambda,
           beta = matrix(0.0, dims["nB"], dims["nL"]),
           pfct = as.double(rep(1.0, dims["nG"])),
           L = new("dCHMsimpl") 
           )

str(ans)
.Call("R_norm_glasso", ans)

ans@beta
ans@a + my

# b) compare with the grplasso pkg
x <- as.matrix(XX)
index = c(rep(1:4, rep(3, 4)), 5, 6)

lambda2 <- lambdamax(x, y = y, index = index, penscale = sqrt, center = FALSE,
                     standardize = TRUE, 
                    model = LinReg())

f2 <- grplasso(as.matrix(X), y, index = c(rep(1:4, rep(3, 4)), 5, 6),
               model = LinReg(), lambda = lambda*2 * length(y),
               center = FALSE, standardize = FALSE,
               control = grpl.control(save.x = T))
f2$coefficients[, 1:5]
ans@beta[, 1:5]
b[14, ]


# c) R implementation
niter <- 200
nbeta <- ncol(X)
ngrp <- length(grp) - 1
eps <- 1e-6

y = y - mean(y)
nlam <- length(lambda)
xm <- colMeans(X)
bb <- matrix(0, nbeta, nlam + 1)
a <- rep(0, nlam + 1)
for (l in 1:nlam){
  b <- matrix(0, niter + 1, nbeta)
  b[1, ] <- bb[, 1]
  lam <- lambda[l]
  r <- as.numeric(y - X %*% b[1, ])
  for (iter in 1:niter){  
    for (j in 1:ngrp){
      idx <- (grp[j] + 1):(grp[j + 1])
      xk <- X[, idx, drop = FALSE]
      rk <- r + xk %*% b[iter, idx]
      xrk <- as.numeric(t(xk) %*% rk)  
      bk <- th_glasso(xrk, lam * sqrt(ncols[j]))
      if (norm(bk) != 0)
        r <- rk - xk %*% bk
      b[iter + 1, idx] <- bk    
    }
    anew <-  - sum(b[iter + 1, ] * xm)
    r <- r + a[l] - anew
    a[l] <- anew
    if (norm(b[iter + 1, ] - b[iter, ])/max(norm(b[iter, ]), 1e-3) < eps)
      break
  }
  bb[, l] <- b[iter + 1, ]
}
bb
a
